/*
Copyright (c) 2010, Advanced Micro Devices
All rights reserved.

Redistribution and use in source and binary forms, with or without modification, are permitted 
provided that the following conditions are met:

Redistributions of source code must retain the above copyright notice, this list of conditions 
and the following disclaimer.

Redistributions in binary form must reproduce the above copyright notice, this list of 
conditions and the following disclaimer in the documentation and/or other materials provided 
with the distribution.

Neither the name of Advanced Micro Devices nor the names of its contributors may be used to endorse 
or promote products derived from this software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR 
IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY 
AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR 
CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR 
CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY 
THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR 
OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE 
POSSIBILITY OF SUCH DAMAGE.
*/

#ifndef __AST_HDR__
#define __AST_HDR__

/*! \file
 *   \brief AST for Pipeline Kit    
 *   \author Benedict R. Gaster 
 *   
 *   \version 0.1
 *   \date March 2010
 */

#include "platform.h"

#ifdef __APPLE__
#include "/boost_1_43_0/boost/variant.hpp"
#include "/boost_1_43_0/boost/variant/recursive_variant.hpp"
#include "/boost_1_43_0/boost/lexical_cast.hpp"
#include "/boost_1_43_0/boost/tuple/tuple.hpp"
#else//Windows
#include <boost/variant.hpp>
#include <boost/variant/recursive_variant.hpp>
#include <boost/lexical_cast.hpp>
#include "boost/tuple/tuple.hpp"
#endif

#include <string>
#include <vector>
#include <list>
#include <utility>
#include <algorithm>
#include <map>

#include <iostream> 

namespace amd {

enum OPERATOR
{
	AST_PLUS,
	AST_MINUS,
	AST_DIVIDE,
	AST_MULTIPLY
};

template <int OpTag> struct ASTBinaryOp;

typedef boost::variant<
  double,                                   
  int,                     
  std::string,             
  boost::recursive_wrapper<ASTBinaryOp<AST_PLUS> >,
  boost::recursive_wrapper<ASTBinaryOp<AST_MINUS> >,
  boost::recursive_wrapper<ASTBinaryOp<AST_DIVIDE> >,  
  boost::recursive_wrapper<ASTBinaryOp<AST_MULTIPLY> > 
  > ASTExpr;

template <int OpTag>
struct ASTBinaryOp
{
  ASTExpr left;
  ASTExpr right;

  ASTBinaryOp(const ASTExpr& lhs, const ASTExpr& rhs)
    : left(lhs), right(rhs)
  {
  }
};

typedef boost::tuple<std::string, ASTExpr> ASTVarDef;
typedef std::map<std::string, ASTVarDef> ASTVarDefMap;

class ExprToString : public boost::static_visitor<std::string>
{
public:
	std::string operator ()(double f) const
	{
		return boost::lexical_cast<std::string>(f);
	}

	std::string operator ()(int i) const
	{
		return boost::lexical_cast<std::string>(i);
	}

	std::string operator ()(std::string var) const
	{
		return boost::lexical_cast<std::string>(var);
	}
  
	template <int OpTag>
	std::string operator ()(ASTBinaryOp<OpTag>& binOp) const
	{
		std::string opTable[] = { "+", "-", "/", "*" };

		std::string str;
		str = "(" + boost::apply_visitor(ExprToString(), binOp.left) + ")"; 

		if (OpTag < AST_PLUS || OpTag > AST_MULTIPLY) {
			str += " UNKNOWN OP ";
		}
		else {
			str += " " + opTable[OpTag] + " " ;
		}

		str += "(" + boost::apply_visitor(ExprToString(), binOp.right) + ")"; 

		return str;
  }
};

class ExprToNameList : public boost::static_visitor<std::list<std::string> >
{
public:
	std::list<std::string> operator ()(double f) const
	{
		return std::list<std::string>();
	}

	std::list<std::string> operator ()(int i) const
	{
		return std::list<std::string>();
	}

	std::list<std::string> operator ()(std::string var) const
	{
		std::list<std::string> list;

		list.push_back(var);

		return list;
	}
  
	template <int OpTag>
	std::list<std::string> operator ()(ASTBinaryOp<OpTag>& binOp) const
	{
		std::list<std::string> list =
			boost::apply_visitor(ExprToNameList(), binOp.left); 

		std::list<std::string> list2 = boost::apply_visitor(ExprToNameList(), binOp.right);
		list.splice(
			list.end(), 
			//boost::apply_visitor(ExprToNameList(), binOp.right));
					list2);

		return list;
  }
};

inline std::list<std::string> 
expr2NameList(ASTExpr& expr)
{
	return boost::apply_visitor(amd::ExprToNameList(), expr);
}

//------------------------------------------------------------------------------

/*! \brief attribute for memory and stage objects
 *
 * Syntax: [name ':'] attrId '=' expression ';'
 */
typedef boost::tuple<std::string, std::string,ASTExpr> ASTAttribute;
typedef std::list<ASTAttribute> ASTAttributeList;

//-------------------------------------------------------------------------------

class ExprToDouble : public boost::static_visitor<double>
{
private:
	const ASTVarDefMap env_;
	const ASTAttributeList attrs_;

	bool findAttr(const std::string name, ASTExpr* expr) const
	{
		bool found = false;

		for (ASTAttributeList::const_iterator i = attrs_.begin(); i != attrs_.end(); i++) {
			if (boost::get<1>(*i).compare(name) == 0) {
				found = true;
				*expr = boost::get<2>(*i);
				break;
			}
		}
		return false;
	}
public:
	ExprToDouble(const ASTVarDefMap& env) : env_(env)
	{
	}

	ExprToDouble(const ASTVarDefMap& env,
				 const ASTAttributeList& attrs) 
		: env_(env), attrs_(attrs)
	{
	}


	double operator ()(double f) const
	{
		return f;
	}

	double operator ()(int i) const
	{
		return static_cast<double>(i);
	}

	double operator ()(const std::string var) const
	{
		double result;
		ASTExpr expr;
		ASTVarDefMap::const_iterator i = env_.find(var);

		if (i != env_.end()) {
			result = boost::apply_visitor(
				ExprToDouble(env_, attrs_), 
				boost::get<1>((*i).second));
		}
		else if (findAttr(var, &expr)) {
			result = boost::apply_visitor(
				ExprToDouble(env_, attrs_), 
				expr);			
		}
		else {
			std::string message = var + " not defined in env";
			throw std::runtime_error( message );
		}

		return result;
	}
  
	template <int OpTag>
	double operator ()(const ASTBinaryOp<OpTag>& binOp) const
	{
		double result;
		double lhs = boost::apply_visitor(ExprToDouble(env_, attrs_), binOp.left); 
		double rhs = boost::apply_visitor(ExprToDouble(env_, attrs_), binOp.right); 

		switch(OpTag) {
		case AST_PLUS:
			result = lhs + rhs;
			break;
		case AST_MINUS:
			result = lhs - rhs;
			break;
		case AST_DIVIDE:
			result = lhs / rhs;
			break;
		case AST_MULTIPLY:
			result = lhs * rhs;
			break;
		default:
			throw std::runtime_error(" unsupported binary operator");
		}
		return result;
  }
};

inline std::string stringify(ASTExpr& expr)
{
	return boost::apply_visitor(amd::ExprToString(), expr);
}

template<typename T>
inline T evaluate(const ASTVarDefMap& env, ASTExpr& expr)
{
	return (T)boost::apply_visitor(amd::ExprToDouble(env), expr);
}

template<typename T>
inline T evaluate(
	const ASTVarDefMap& env, 
	const ASTAttributeList& attr,
	ASTExpr& expr)
{
	return (T)boost::apply_visitor(amd::ExprToDouble(env, attr), expr);
}
//----------------------------------------------------------------------------------

inline std::string stringify(ASTVarDefMap& vardefs)
{
	std::string str("\n");

	for (ASTVarDefMap::iterator i = vardefs.begin(); i != vardefs.end(); i++) {
		str +=  
			boost::get<0>((*i).second) + " " +
			(*i).first + " = " +
			stringify(boost::get<1>((*i).second)) + 
			";\n";
	}

	str += "\n";

	return str;
}

//----------------------------------------------------------------------------------

inline std::string stringify(ASTAttribute& attr)
{
	std::string str;

	if (boost::get<0>(attr).length() > 0) {
		str = boost::get<0>(attr) + " : \n\t\t" ;
	}

	str += boost::get<1>(attr) + " = " ;
	str += boost::apply_visitor(ExprToString(), boost::get<2>(attr)) + ";";

	return str;
}

inline std::string stringify(ASTAttributeList& attrs)
{
	std::string str;

	for (ASTAttributeList::iterator i = attrs.begin(); i != attrs.end(); i++) {
		str += "\t" + stringify(*i) + "\n";
	}

	return str;
}

//----------------------------------------------------------------------------------

enum ADDRESS_SPACE
{
	AST_GLOBAL,
	AST_LOCAL,
	AST_UNIFORM
};

struct ASTBuffer
{
	ADDRESS_SPACE aspace_;
	std::string type_;
	std::string name_;
	ASTExpr size_;

	// attributes
	ASTAttributeList attrs_;
};

struct ASTImage2D
{
	ADDRESS_SPACE aspace_;
	std::string type_;
	std::string name_;
	ASTExpr sizeX_;
	ASTExpr sizeY_;

	// attributes
	ASTAttributeList attrs_;
};

struct ASTImage3D
{
	ADDRESS_SPACE aspace_;
	std::string type_;
	std::string name_;
	ASTExpr sizeX_;
	ASTExpr sizeY_;
	ASTExpr sizeZ_;

	// attributes
	ASTAttributeList attrs_;
};

typedef boost::variant<
	ASTBuffer,
	ASTImage2D,
	ASTImage3D
	> ASTMemory;

template <typename T>
class Maybe
{
private:
	class MaybeToBool : public boost::static_visitor<bool>
	{
	public:
		bool operator ()(T& value) const
		{
			return true;
		}

		bool operator ()(boost::tuple<std::string>& str) const
		{
			return false;
		}
	};

	class MaybeToT : public boost::static_visitor<T>
	{
	public:
		T operator ()(T& value) const
		{
			return value;
		}

		T operator ()(boost::tuple<std::string>& str) const
		{
			throw std::runtime_error("cannot convert a bool to a T");
			return T();
		}
	};

public:
	typedef boost::variant<T, boost::tuple<std::string> > MaybeType;

	static bool toBool(typename Maybe<T>::MaybeType& maybe)
	{
		return boost::apply_visitor(MaybeToBool(), maybe);
	}

	static T toJust(typename Maybe<T>::MaybeType& maybe)
	{
		return boost::apply_visitor(MaybeToT(), maybe);
	}
};


class MemoryToString : public boost::static_visitor<std::string>
{
public:
	std::string operator ()(ASTBuffer& buffer) const
	{
		std::string str("Buffer " + buffer.name_ + "(");

		switch(buffer.aspace_) {
		case AST_GLOBAL:		
			str += "__global, ";
			break;
		case AST_LOCAL:
			str += "__local, ";
			break;
		case AST_UNIFORM:
			str += "__uniform, ";
			break;
		default:
			str += "UNKNOWN ADDRESS SPACE, ";
			break;
		}

		str += buffer.type_ + ", ";

		str += boost::apply_visitor(ExprToString(), buffer.size_);

		str += ")";

		// attributes
		str += "\n{\n";
		str += stringify(buffer.attrs_);
		str += "\n}\n";

		return str;
	}

	std::string operator ()(ASTImage2D& image) const
	{
		std::string str("Image2D " + image.name_ + "(");

		switch(image.aspace_) {
		case AST_GLOBAL:		
			str += "__global, ";
			break;
		case AST_LOCAL:
			str += "__local, ";
			break;
		case AST_UNIFORM:
			str += "__uniform, ";
			break;
		default:
			str += "UNKNOWN ADDRESS SPACE, ";
			break;
		}

		str += image.type_ + ", ";

		str += boost::apply_visitor(ExprToString(), image.sizeX_);
		str += ", ";
		str += boost::apply_visitor(ExprToString(), image.sizeY_);

		str += ")";

		// attributes
		str += "\n{\n";
		str += stringify(image.attrs_);
		str += "\n}\n";

		return str;
	}

	std::string operator ()(ASTImage3D& image) const
	{
		std::string str("Image3D " + image.name_ + "(");

		switch(image.aspace_) {
		case AST_GLOBAL:		
			str += "__global, ";
			break;
		case AST_LOCAL:
			str += "__local, ";
			break;
		case AST_UNIFORM:
			str += "__uniform, ";
			break;
		default:
			str += "UNKNOWN ADDRESS SPACE, ";
			break;
		}

		str += image.type_ + ", ";

		str += boost::apply_visitor(ExprToString(), image.sizeX_);
		str += ", ";
		str += boost::apply_visitor(ExprToString(), image.sizeY_);
		str += ", ";
		str += boost::apply_visitor(ExprToString(), image.sizeZ_);

		str += ")";

		// attributes
		str += "\n{\n";
		str += stringify(image.attrs_);
		str += "\n}\n";

		return str;
	}
};

typedef std::list<ASTMemory> ASTMemoryList;

inline std::string stringify(ASTMemoryList& mems)
{
	std::string str;

	for (ASTMemoryList::iterator i = mems.begin(); i != mems.end(); i++) {
		str += boost::apply_visitor( MemoryToString(), *i) + "\n";
	}

	return str;
}

//----------------------------------------------------------------------------------

class ASTUniform
{
public:
	std::string name_;
	std::string type_;

	// attributes
	ASTAttributeList attrs_;
};

typedef std::list<ASTUniform> ASTUniformList;

inline std::string stringify(ASTUniformList& uniforms)
{
	std::string str;

	for (ASTUniformList::iterator i = uniforms.begin(); i != uniforms.end();) {
		str += "Uniform (" + (*i).type_ + " " + (*i).name_ + ")";

		// attributes
		str += "\n{\n";
		str += stringify((*i).attrs_);
		str += "\n}\n";

		i++;
		if (i != uniforms.end()) {
			str += " ";
		}
	}

	return str;
}

//----------------------------------------------------------------------------------

enum CODE_SECTION
{
	AST_CODECPP,
	AST_CODECL
};

class ASTCode
{
public:
	CODE_SECTION type_;
	std::string code_;
	// attributes
	ASTAttributeList attrs_;
};

typedef std::list<ASTCode> ASTCodeList;

inline std::string stringify(ASTCodeList& codes)
{
	std::string str;

	for (ASTCodeList::iterator i = codes.begin(); i != codes.end(); i++) {
		ASTCode code = *i;

		str += code.type_ == AST_CODECPP ? "CodeCPP" : "CodeCL";
		str += "\n{\n";
		str += code.code_;
		str += "\n}\n\n";
	}

	return str;
}


//----------------------------------------------------------------------------------

typedef std::list<std::string> ASTTypeList;

typedef std::string TypeName;
typedef std::string VarName;
typedef std::string IoClass;
typedef boost::tuple<ADDRESS_SPACE, TypeName, VarName, int, IoClass> ASTTypeVarElement;
typedef std::list<ASTTypeVarElement > ASTTypeVarList;

inline std::string stringify(ASTTypeList& types)
{
	std::string str;

	for (ASTTypeList::iterator i = types.begin(); i != types.end();) {
		str += *(i);
		i++;
		if (i != types.end()) {
			str += ", ";
		}
	}

	return str;
}

inline std::string stringify(ASTTypeVarList& types)
{
	std::string str;

	for (ASTTypeVarList::iterator i = types.begin(); i != types.end();) {
		str += boost::get<0>(*(i)) + " " + boost::get<1>(*(i)) + " " + boost::get<2>(*(i)) ;
		i++;
		if (i != types.end()) {
			str += ", ";
		}
	}

	return str;
}

class ASTKernel
{
public:
	std::string name_;

#if 1
	ASTTypeVarList arguments_;
#else
	ASTTypeVarList inputs_;
	ASTTypeVarList outputs_;
#endif

	// attributes
	ASTAttributeList attrs_;
};

typedef std::list<ASTKernel> ASTKernelList;

inline std::string stringify(ASTKernel& kernel)
{
	std::string str("Kernel " + kernel.name_ + " : ");
#if 1
	str += " (" + stringify( kernel.arguments_ ) + ")";
#else
	str += " (" + stringify(kernel.inputs_) + ")";	
	str += " ->";
	str += " (" + stringify(kernel.outputs_) + ")";
#endif

	// attributes
	str += "\n{\n";
	str += stringify(kernel.attrs_);
	str += "\n}\n";

	return str;
}

inline std::string stringify(ASTKernelList& kernels)
{
	std::string str;

	for (ASTKernelList::iterator i = kernels.begin(); i != kernels.end(); i++) {
		str += stringify(*i);
	}

	return str;
}

class ASTStage
{
public:
	std::string name_;
	std::string kernelName_;

#if 0
	ASTTypeList readonlyArgs_;
	ASTTypeList readwriteArgs_;
#else
	ASTTypeList args_;
#endif

	// attributes
	ASTAttributeList attrs_;
};

typedef std::list<ASTStage> ASTStageList;

inline std::string stringify(ASTStage& stage)
{
	std::string str("Stage " + stage.name_ + " : " + stage.kernelName_);
#if 1
	str += " (" + stringify( stage.args_ ) + ")";
#else
	str += " (" + stringify(stage.readonlyArgs_) + ")";	
	str += " ->";
	str += " (" + stringify(stage.readwriteArgs_) + ")";
#endif

	// attributes
	str += "\n{\n";
	str += stringify(stage.attrs_);
	str += "\n}\n";

	return str;
}

inline std::string stringify(ASTStageList& stages)
{
	std::string str;

	for (ASTStageList::iterator i = stages.begin(); i != stages.end(); i++) {
		str += stringify(*i);
	}

	return str;
}

//----------------------------------------------------------------------------------

typedef boost::tuple<std::string, std::string> ASTGraphEdge; 
typedef std::list<ASTGraphEdge> ASTGraphEdgeList;

class ASTGraph
{
public:
	std::string name_;
	ASTGraphEdgeList edges_;

	// attributes
	ASTAttributeList attrs_;
};

inline std::string stringify(ASTGraph& graph)
{
	std::string str;

	str = "Graph " + graph.name_;
	str += "\n{\n";
	
	for (ASTGraphEdgeList::iterator i = graph.edges_.begin(); i != graph.edges_.end();) {
		str += "\t" + boost::get<0>(*(i)) + ", " + boost::get<1>(*(i)) + ";";
		i++;
		if (i != graph.edges_.end()) {
			str += "\n";
		}
	}
	
	str += "\n}";

	// attributes
	str += "\n{\n";
	str += stringify(graph.attrs_);
	str += "\n}\n";

	return str;
}

//----------------------------------------------------------------------------------

class ASTPipeline
{
public:
	std::string name_;
	ASTCodeList codes_;
	ASTVarDefMap environment_;
	ASTMemoryList memorys_;
	ASTUniformList uniforms_;
	ASTKernelList kernels_;
	ASTStageList stages_;
	ASTGraph graph_;
	
	// attributes
	ASTAttributeList attrs_;
};

inline std::string stringify(ASTPipeline& pipeline)
{
	std::string str;

	str = "Pipeline " + pipeline.name_;

	// attributes
	str += "\n{\n";
	str += stringify(pipeline.attrs_);
	str += "\n}\n";

	str += "\n{\n";

	str += stringify(pipeline.codes_);

	str += stringify(pipeline.environment_);

	str += stringify(pipeline.memorys_);

	str += "\n\n";

	str += stringify(pipeline.uniforms_);

	str += "\n\n";

	str += stringify(pipeline.kernels_);

	str += "\n\n";

	str += stringify(pipeline.stages_);

	str += "\n\n";

	str += stringify(pipeline.graph_);

	return str;
}

}; // namespace amd
#endif
